package HandshakeJoin;

import java.util.List;
import java.util.ArrayList;
import com.google.common.collect.Lists;
import com.google.common.primitives.Ints;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class CassandraClient implements DBClient
{
    transient private static final Logger logger = LoggerFactory
            .getLogger(CassandraClient.class);

    transient static final private String connectIP = "127.0.0.1";
    transient static final private String keyspace = "handshakejoin";

    public static void initGlobal() {
        Cluster cluster = Cluster.builder().addContactPoint(connectIP).build();
        Session session = cluster.connect();
        String cql;

        cql = "DROP KEYSPACE " + keyspace + ";";
        logger.info("CQL query: " + cql);
        session.execute(cql);

        cql = "CREATE KEYSPACE "
                + keyspace
                + " WITH replication = {\'class\': \'SimpleStrategy\', \'replication_factor\': 1 };";
        logger.info("CQL query: " + cql);
        session.execute(cql);

        session.shutdown();
        cluster.shutdown();
    }

    private String tableName;
    private int numColumns;
    private Cluster cluster;
    private Session session;
    private long seq = 0;

    public CassandraClient(final String tableName, final int numColumns) {
        this.tableName = tableName;
        this.numColumns = numColumns;
        cluster = Cluster.builder().addContactPoint(connectIP).build();
        session = cluster.connect();

        query("USE " + keyspace + ";");
    }

    public boolean init() {
        StringBuilder sb = new StringBuilder();
        sb.append("CREATE TABLE " + tableName + " (");
        sb.append("seq bigint PRIMARY KEY");
        for (int i = 1; i <= numColumns; ++i) {
            sb.append(", col_" + i + " int");
        }
        sb.append(", timestamp bigint);");

        query(sb.toString());

        query("CREATE INDEX ON " + tableName + " (timestamp);");

        query("CREATE INDEX ON " + tableName + " (col_1);");

        return true;
    }

    public void close() {
        session.shutdown();
        cluster.shutdown();
    }

    public boolean materialize(Snapshot snapshot) {
        for (ArrayList<Integer> record : snapshot.getRecords()) {
            insert(record, snapshot.getTimeStamp());
        }

        return true;
    }

    public List<ArrayList<Integer>> retrieveByTimestamp(long ts) {
        return convert(selectByTimestamp(ts));
    }

    public List<ArrayList<Integer>> retrieveByTimestamp(long tsLower,
            long tsUpper) {
        /* Not supported by Cassandra. */
        return null;
    }

    public List<ArrayList<Integer>> retrieveByJoinKey(int val) {
        return convert(selectByJoinKey(val));
    }

    public List<ArrayList<Integer>> retrieveByJoinKey(int valLower, int valUpper) {
        /* Not supported by Cassandra. */
        return null;
    }

    private ResultSet query(final String cql) {
        // logger.info("CQL query: " + cql);

        return session.execute(cql);
    }

    private void insert(final ArrayList<Integer> record, final long timestamp) {
        StringBuilder sb = new StringBuilder();
        sb.append("INSERT INTO " + tableName + " (");
        sb.append("seq");
        for (int i = 1; i <= numColumns; ++i) {
            sb.append(", col_" + i);
        }
        sb.append(", timestamp) VALUES (");
        sb.append(Long.toString(++seq));
        for (int i = 0; i < numColumns; ++i) {
            sb.append(", " + record.get(i));
        }
        sb.append(", " + timestamp + ");");

        query(sb.toString());
    }

    private ResultSet selectByTimestamp(final long timestamp) {
        StringBuilder sb = new StringBuilder();
        sb.append("SELECT col_1");
        for (int i = 2; i <= numColumns; ++i) {
            sb.append(", col_" + i);
        }
        sb.append(" FROM " + tableName);
        sb.append(" WHERE timestamp = " + timestamp + ";");

        return query(sb.toString());
    }

    private ResultSet selectByJoinKey(final int val) {
        StringBuilder sb = new StringBuilder();
        sb.append("SELECT col_1");
        for (int i = 2; i <= numColumns; ++i) {
            sb.append(", col_" + i);
        }
        sb.append(" FROM " + tableName);
        sb.append(" WHERE col_1 = " + val + ";");

        return query(sb.toString());
    }

    private List<ArrayList<Integer>> convert(final ResultSet rs) {
        List<ArrayList<Integer>> records = Lists.newLinkedList();

        for (Row row : rs) {
            ArrayList<Integer> record = Lists
                    .newArrayListWithCapacity(numColumns);

            for (int i = 0; i < numColumns; ++i) {
                record.add(row.getInt(i));
            }

            records.add(record);
        }

        return records;
    }

    /**
     * Test client.
     */
    public static void main(String[] args) {
        CassandraClient.initGlobal();

        final int numColumns = 3;
        DBClient db = new CassandraClient("test", numColumns);
        db.init();

        int[][] data = { { 1, 2, 3 }, { 1, 3, 4 }, { 2, 5, 3 }, { 3, 2, 6 },
                { 1, 3, 2 }, { 4, 5, 1 } };
        List<ArrayList<Integer>> records = Lists.newLinkedList();
        for (int i = 0; i < data.length; ++i) {
            records.add(Lists.newArrayList(Ints.asList(data[i])));
        }
        Snapshot snapshot = new Snapshot("test", records);

        db.materialize(snapshot);

        records = db.retrieveByJoinKey(1);
        for (ArrayList<Integer> record : records) {
            System.out.println(record);
        }

        db.close();
    }
}